# database.py
import sqlite3
import csv
import os
from datetime import datetime, date

class Database:
    def __init__(self, db_path="accounting.db"):
        self.db_path = db_path
        self.init_db()
    
    def init_db(self):
        """初始化数据库"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建交易表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS transactions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                type TEXT NOT NULL,
                amount REAL NOT NULL,
                main_category TEXT NOT NULL,
                sub_category TEXT NOT NULL,
                item_detail TEXT,
                note TEXT,
                date TEXT NOT NULL,
                time TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建预算表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS budgets (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                category TEXT NOT NULL,
                amount REAL NOT NULL,
                period TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建预算预警设置表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS budget_alerts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                threshold INTEGER NOT NULL DEFAULT 80,
                is_enabled BOOLEAN NOT NULL DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 检查是否需要插入默认预算预警设置
        try:
            cursor.execute('SELECT COUNT(*) FROM budget_alerts')
            if cursor.fetchone()[0] == 0:
                cursor.execute('''
                    INSERT INTO budget_alerts (threshold, is_enabled)
                    VALUES (80, 0)
                ''')
        except:
            # 如果表结构不匹配，跳过插入
            pass
        
        conn.commit()
        conn.close()
    
    def add_transaction(self, transaction_type, amount, main_category, sub_category, item_detail="", note=""):
        """添加交易记录"""
        # 确保金额保留两位小数
        amount = round(amount, 2)
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO transactions (type, amount, main_category, sub_category, item_detail, note, date, time)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        ''', (transaction_type, amount, main_category, sub_category, item_detail, note, 
              datetime.now().strftime("%Y-%m-%d"), datetime.now().strftime("%H:%M:%S")))
        
        conn.commit()
        conn.close()
    
    def get_transactions(self, start_date=None, end_date=None):
        """获取交易记录"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        query = "SELECT * FROM transactions WHERE 1=1"
        params = []
        
        if start_date:
            query += " AND date >= ?"
            params.append(start_date)
        if end_date:
            query += " AND date <= ?"
            params.append(end_date)
        
        cursor.execute(query, params)
        transactions = cursor.fetchall()
        
        conn.close()
        
        # 转换为字典列表
        return [
            {
                'id': row[0],
                'type': row[1],
                'amount': float(row[2]),
                'main_category': row[3],
                'sub_category': row[4],
                'item_detail': row[5],
                'note': row[6],
                'date': row[7],
                'time': row[8]
            }
            for row in transactions
        ]
    
    def get_transaction_by_details(self, date, time, amount):
        """根据日期、时间和金额查找交易记录"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT * FROM transactions 
            WHERE date = ? AND time = ? AND amount = ?
        ''', (date, time, amount))
        
        result = cursor.fetchone()
        conn.close()
        
        if result:
            return {
                'id': result[0],
                'type': result[1],
                'amount': float(result[2]),
                'main_category': result[3],
                'sub_category': result[4],
                'item_detail': result[5],
                'note': result[6],
                'date': result[7],
                'time': result[8]
            }
        return None
    
    def delete_transaction(self, transaction_id):
        """删除交易记录"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('DELETE FROM transactions WHERE id = ?', (transaction_id,))
        conn.commit()
        conn.close()
        
        return cursor.rowcount > 0
    
    def export_to_csv(self, file_path, start_date=None, end_date=None):
        """导出交易记录到CSV文件"""
        try:
            # 获取交易记录
            transactions = self.get_transactions(start_date, end_date)
            
            if not transactions:
                return False
            
            # 写入CSV文件
            with open(file_path, 'w', newline='', encoding='utf-8') as csvfile:
                writer = csv.writer(csvfile)
                
                # 写入标题行
                writer.writerow(['日期', '时间', '类型', '金额', '大类', '分类', '具体项目', '备注'])
                
                # 写入数据行
                for transaction in transactions:
                    # 金额保留两位小数
                    amount = f"{transaction['amount']:.2f}"
                    writer.writerow([
                        transaction['date'],
                        transaction['time'],
                        '收入' if transaction['type'] == 'income' else '支出',
                        amount,
                        transaction['main_category'],
                        transaction['sub_category'],
                        transaction['item_detail'] or '',
                        transaction['note'] or ''
                    ])
            
            return True
            
        except Exception as e:
            print(f"导出失败: {e}")
            return False
    
    def export_summary_to_csv(self, file_path, period='month', start_date=None, end_date=None):
        """导出汇总数据到CSV文件"""
        try:
            # 获取汇总数据
            summary_data = self.get_summary_data(period, start_date, end_date)
            
            if not summary_data:
                return False
            
            # 写入CSV文件
            with open(file_path, 'w', newline='', encoding='utf-8') as csvfile:
                writer = csv.writer(csvfile)
                
                # 写入标题行
                if period == 'day':
                    writer.writerow(['日期', '总收入', '总支出', '结余', '交易笔数'])
                elif period == 'week':
                    writer.writerow(['周次', '总收入', '总支出', '结余', '交易笔数'])
                elif period == 'quarter':
                    writer.writerow(['季度', '总收入', '总支出', '结余', '交易笔数'])
                elif period == 'year':
                    writer.writerow(['年份', '总收入', '总支出', '结余', '交易笔数'])
                else:  # month
                    writer.writerow(['月份', '总收入', '总支出', '结余', '交易笔数'])
                
                # 写入数据行
                for data in summary_data:
                    # 格式化季度名称
                    if period == 'quarter':
                        quarter_name = self.format_quarter(data['period'])
                    else:
                        quarter_name = data['period']
                    
                    # 金额保留两位小数
                    income = f"{data['total_income']:.2f}"
                    expense = f"{data['total_expense']:.2f}"
                    balance = f"{data['balance']:.2f}" if 'balance' in data else f"{(data['total_income'] - data['total_expense']):.2f}"
                    
                    writer.writerow([
                        quarter_name,
                        income,
                        expense,
                        balance,
                        data['count']
                    ])
            
            return True
            
        except Exception as e:
            print(f"导出失败: {e}")
            return False
    
    def get_summary_data(self, period='month', start_date=None, end_date=None):
        """获取汇总数据"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 获取所有交易记录，支持日期范围筛选
        query = 'SELECT date, type, amount FROM transactions WHERE 1=1'
        params = []
        
        if start_date:
            query += " AND date >= ?"
            params.append(start_date)
        if end_date:
            query += " AND date <= ?"
            params.append(end_date)
            
        query += ' ORDER BY date'
        
        cursor.execute(query, params)
        records = cursor.fetchall()
        
        conn.close()
        
        if not records:
            return []
        
        # 按照指定周期分组计算
        summary_data = []
        current_group = None
        
        for record in records:
            date_str = record[0]
            year, month, day = map(int, date_str.split('-'))
            amount = record[2]
            
            # 根据周期确定分组键
            if period == 'day':
                group_key = date_str
            elif period == 'week':
                # 计算周次
                week_start = datetime(year, month, day) - datetime(year, month, day).replace(day=1).replace(hour=0, minute=0, second=0, microsecond=0)
                week_start = week_start.replace(day=1)
                week_num = (datetime(year, month, day) - week_start).days // 7 + 1
                group_key = f"{year}-第{week_num}周"
            elif period == 'month':
                group_key = f"{year}-{month:02d}"
            elif period == 'quarter':
                quarter = (month - 1) // 3 + 1
                group_key = f"{year}-Q{quarter}"
            else:  # year
                group_key = str(year)
            
            # 处理当前组
            if current_group is None or current_group['period'] != group_key:
                if current_group:
                    summary_data.append(current_group)
                
                current_group = {
                    'period': group_key,
                    'total_income': 0,
                    'total_expense': 0,
                    'count': 0
                }
            
            # 更新当前组的统计
            if record[1] == 'income':
                current_group['total_income'] += amount
            else:
                current_group['total_expense'] += amount
            
            current_group['count'] += 1
        
        # 添加最后一个组
        if current_group:
            # 计算结余
            current_group['balance'] = current_group['total_income'] - current_group['total_expense']
            summary_data.append(current_group)
        
        # 按照时间顺序排序
        summary_data.sort(key=lambda x: x['period'])
        
        return summary_data
    def format_quarter(self, quarter_str):
        """格式化季度显示"""
        try:
            if '-' in quarter_str and quarter_str.startswith('2'):  # 看起来像年份-季度格式
                year, quarter = quarter_str.split('-')
                quarter = int(quarter[1]) if quarter.startswith('Q') else int(quarter)  # Q1, Q2, Q3, Q4 或 1, 2, 3, 4
                
                # 映射季度到中文季节
                season_map = {
                    1: '春',
                    2: '夏',
                    3: '秋',
                    4: '冬'
                }
                
                season = season_map.get(quarter, '')
                return f"{year}年{season}"
            else:
                return quarter_str
        except:
            return quarter_str
    
    def get_available_years(self):
        """获取有数据的年份列表"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT DISTINCT strftime('%Y', date) as year 
            FROM transactions 
            ORDER BY year
        ''')
        
        years = [int(row[0]) for row in cursor.fetchall()]
        conn.close()
        
        return years
    
    def set_budget(self, category, amount, period):
        """设置预算"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 检查是否已存在该分类和周期的预算
        cursor.execute('''
            SELECT id FROM budgets 
            WHERE category = ? AND period = ?
        ''', (category, period))
        
        result = cursor.fetchone()
        if result:
            # 更新现有预算
            cursor.execute('''
                UPDATE budgets 
                SET amount = ? 
                WHERE category = ? AND period = ?
            ''', (amount, category, period))
        else:
            # 插入新预算
            cursor.execute('''
                INSERT INTO budgets (category, amount, period)
                VALUES (?, ?, ?)
            ''', (category, amount, period))
        
        conn.commit()
        conn.close()
    
    def get_budgets(self, period):
        """获取预算"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT category, amount FROM budgets 
            WHERE period = ?
        ''', (period,))
        
        results = cursor.fetchall()
        conn.close()
        
        return [{'main_category': row[0], 'amount': row[1]} for row in results]
    
    def get_budget_status(self, year, month):
        """获取预算执行状态"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 获取该月的所有支出
        cursor.execute('''
            SELECT main_category, SUM(amount) as total_amount
            FROM transactions 
            WHERE type = 'expense' 
            AND strftime('%Y', date) = ? 
            AND strftime('%m', date) = ?
            GROUP BY main_category
        ''', (str(year), f"{int(month):02d}"))
        
        actual_expenses = {row[0]: row[1] for row in cursor.fetchall()}
        
        # 获取月度预算
        cursor.execute('''
            SELECT category, amount FROM budgets 
            WHERE period = 'monthly'
        ''')
        
        budgets = {row[0]: row[1] for row in cursor.fetchall()}
        
        conn.close()
        
        # 计算每个分类的预算状态
        budget_status = []
        all_categories = set(list(actual_expenses.keys()) + list(budgets.keys()))
        
        for category in all_categories:
            budget_amount = budgets.get(category, 0)
            actual_amount = actual_expenses.get(category, 0)
            remaining = budget_amount - actual_amount
            usage_rate = (actual_amount / budget_amount * 100) if budget_amount > 0 else 0
            is_over_budget = actual_amount > budget_amount if budget_amount > 0 else False
            
            budget_status.append({
                'category': category,
                'budget': budget_amount,
                'actual': actual_amount,
                'remaining': remaining,
                'usage_rate': usage_rate,
                'is_over_budget': is_over_budget
            })
        
        return budget_status
    
    def set_budget_alert(self, threshold):
        """设置预算预警阈值"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 检查是否已存在预算预警设置
        try:
            cursor.execute('SELECT COUNT(*) FROM budget_alerts')
            if cursor.fetchone()[0] == 0:
                # 插入新的预算预警设置
                cursor.execute('''
                    INSERT INTO budget_alerts (threshold, is_enabled)
                    VALUES (?, ?)
                ''', (threshold, 1))
            else:
                # 更新现有的预算预警设置
                cursor.execute('''
                    UPDATE budget_alerts 
                    SET threshold = ?, is_enabled = 1
                ''', (threshold,))
        except:
            # 如果表结构有问题，跳过操作
            pass
        
        conn.commit()
        conn.close()
    
    def get_budget_alert(self):
        """获取预算预警设置"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            cursor.execute('''
                SELECT threshold, is_enabled FROM budget_alerts 
                ORDER BY id DESC LIMIT 1
            ''')
            
            result = cursor.fetchone()
            if result:
                return {
                    'threshold': result[0],
                    'is_enabled': bool(result[1])
                }
        except:
            # 如果查询失败，返回默认值
            pass
        finally:
            conn.close()
        
        return {
            'threshold': 80,
            'is_enabled': False
        }